package com.biubiux2.handler;

import com.biubiux2.message.MessageType;
import com.biubiux2.security.DESUtil;
import com.biubiux2.security.KeyNegotiationContext;
import com.biubiux2.security.RSAUtil;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.MessageToMessageCodec;
import io.netty.handler.codec.MessageToMessageDecoder;
import io.netty.handler.codec.MessageToMessageEncoder;
import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import java.nio.charset.StandardCharsets;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.LinkedList;
import java.util.List;

/**
 * @author noexcs
 * @since 8/10/2022 7:14 PM
 */
@Slf4j
public class ClientEncryptHandler extends MessageToMessageCodec<ByteBuf, ByteBuf> {

    private Cipher encryptCipher;

    private Cipher decryptCipher;

    private boolean negotiated = false;

    private final LinkedList<ByteBuf> msgWhilePrenegotiated = new LinkedList<>();

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        super.channelActive(ctx);
        KeyNegotiationContext negotiationContext = new KeyNegotiationContext();
        ctx.channel().attr(KeyNegotiationContext.KEY_NEGOTIATION_CONTEXT_ATTRIBUTE_KEY).set(negotiationContext);
        negotiationContext.negotiate(ctx);
    }

    /**
     * @see MessageToMessageEncoder#write(io.netty.channel.ChannelHandlerContext, java.lang.Object, io.netty.channel.ChannelPromise) 会进行release
     */
    @Override
    protected void encode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {
        if (negotiated) {
            byte[] bytes = new byte[msg.readableBytes()];
            msg.readBytes(bytes);
            byte[] encrypted = encryptCipher.doFinal(bytes);

            ByteBuf encryptedBytes = PooledByteBufAllocator.DEFAULT.buffer(encrypted.length);
            encryptedBytes.writeBytes(encrypted);
            out.add(encryptedBytes);
        } else {
            msgWhilePrenegotiated.addLast(msg.retainedDuplicate());
        }
    }

    /**
     * @see MessageToMessageDecoder#channelRead(io.netty.channel.ChannelHandlerContext, java.lang.Object) 会进行release
     */
    @Override
    protected void decode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {
        if (negotiated) {
            byte[] bytes = new byte[msg.readableBytes()];
            msg.readBytes(bytes);
            byte[] decrypted = decryptCipher.doFinal(bytes);

            ByteBuf decryptedBytes = PooledByteBufAllocator.DEFAULT.buffer(decrypted.length);
            decryptedBytes.writeBytes(decrypted);
            out.add(decryptedBytes);
        } else {
            // 密钥未协商阶段
            int type = msg.readInt();
            if (type == MessageType.SERVER_SECRET_KEY) {
                // 客户端收到对称密钥
                byte[] bytes = new byte[msg.readableBytes()];
                msg.readBytes(bytes);

                // 解密对称密钥
                PrivateKey privateKey = ctx.channel().attr(KeyNegotiationContext.KEY_NEGOTIATION_CONTEXT_ATTRIBUTE_KEY).get().privateKey;
                String decryptSecretKey = RSAUtil.decrypt(new String(bytes, StandardCharsets.UTF_8), privateKey);
                SecretKey secretKey = DESUtil.convertStringToKey(decryptSecretKey);
                encryptCipher = DESUtil.getEncryptCipher(secretKey);
                decryptCipher = DESUtil.getDecryptCipher(secretKey);

                // 标记协商成功
                negotiated = true;
                log.debug("密钥协商成功");

                // 清除未协商前的消息
                while (!msgWhilePrenegotiated.isEmpty()) {
                    ctx.channel().writeAndFlush(msgWhilePrenegotiated.removeFirst());
                }
            } else {
                log.debug("密钥未协商");
            }
        }
    }
}
